# -*- coding: utf-8 -*-
# time: 2024/5/8 16:36
# file: doc_Loaders_chain.py
# author: hanson
"""
主要演示文档检索和AI对话,这个非常有用，运营支持只需要整理多份 百度FAQ.txt 文档即可以啦
https://cloud.baidu.com/qianfandev/topic/267671
构建本地知识库问答机器人
安装 langchain 相关依赖包
pip install langchain
pip install langchain-community
pip install langchain-core
pip install langchain-experimental
pip install “langserve[all]”
pip install langchain-cli
pip install langsmith
pip install sentence-transformers
pip install chromadb
"""
from langchain.chains.retrieval_qa.base import RetrievalQA
from langchain_community.document_loaders import  DirectoryLoader

from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.llms.ollama import Ollama
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import PromptTemplate


from langchain_text_splitters import CharacterTextSplitter
# 加载文件夹中的所有.md类型的文件 这里需要绝对路径？？
""""
  Document loaders 文档加载器
重点包括了 txt（TextLoader）、csv（CSVLoader），html（UnstructuredHTMLLoader），json（JSONLoader），markdown（UnstructuredMarkdownLoader）以及 pdf（因为 pdf 的格式比较复杂，
提供了 PyPDFLoader、MathpixPDFLoader、UnstructuredPDFLoader，PyMuPDF 等多种形式的加载引擎）几种常用格式的内容解析。
"""
loader = DirectoryLoader('F:/workspace/py_project/intellect/llm/data',glob="**/*.md", use_multithreading=True)
# 将数据转成 Document 对象，每个文件会作为一个 Document
documents = loader.load()
# 初始化加载器
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
# 切割加载的 document
split_docs = text_splitter.split_documents(documents)

"""
git clone https://www.modelscope.cn/thomas/text2vec-base-chinese.git
中文转向量下载库
"""
# 初始化 embeddings 对象
embeddings = HuggingFaceEmbeddings(model_name='E:/opt/AI/text2vec-base-chinese')
# 将 document 计算 embedding 向量信息并临时存入 Chroma 向量数据库，用于后续匹配查询
# 将 document 计算 embedding 向量信息并临时存入 Chroma 向量数据库，用于后续匹配查询
docsearch = Chroma.from_documents(split_docs, embeddings, persist_directory="F:/temp/cache")
docsearch.persist()
# 加载数据
docsearch = Chroma(persist_directory="F:/temp/cache", embedding_function=embeddings)

# 创建问答对象 llm模型
llm = Ollama(model="gemma:qwen2.5:1.5b")
prompt_template = """基于以下已知信息，简洁和专业的来回答用户的问题。
如果无法从中得到答案，请说 "根据已知信息无法回答该问题" 或 "没有提供足够的相关信息"，不允许在答案中添加编造成分，答案请使用代码。
已知内容:
{context}
问题:
{question}"""
prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
chain_type_kwargs = {"prompt": prompt}
qa = RetrievalQA.from_chain_type(llm=llm, retriever=docsearch.as_retriever(), chain_type="stuff",
                                 chain_type_kwargs=chain_type_kwargs, return_source_documents=True)

# 进行问答
query = "我的认证信息会泄露吗？"
res = qa({"query": query})
answer, docs = res['result'], res['source_documents']
print("\n\n> 问题:")
print(query)
print("\n> 回答:")
print(answer)
for document in docs:
    print("\n> " + document.metadata["source"] + ":")

